package mycollections;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	static final boolean RED = false;
	static final boolean BLACK = true;
	SimpleEntry<K, V> root = null;
	int size;
	private Comparator<K> comparator;
	
	public MyTreeMap() {
	}

	public MyTreeMap(Comparator<K> comparator) {
		this.comparator = comparator;
	}

	public Iterator<MyMap.Entry<K, V>> entryIterator() {
		return new EntryIterator(getFirstEntry());
	}

	public boolean containsKey(K key) {
		return getEntry(key) != null;
	}

	public boolean containsValue(V value) {
		for (SimpleEntry<K, V> e = getFirstEntry(); e != null; e = getSucc(e))
			if (equalsOrNull(value, e.value))
				return true;
		return false;
	}

	public int size() {
		return size;
	}

	public V put(K key, V value) {
		SimpleEntry<K, V> e = root;
		if (e == null) {
			compare(key, key);
			root = new SimpleEntry<K, V>(key, value, null);
			size = 1;
			return null;
		}
		int cmp = 0;
		SimpleEntry<K, V> parent = null;
		do {
			parent = e;
			cmp = compare(key, e.key);
			if (cmp < 0)
				e = e.left;
			else if (cmp > 0)
				e = e.right;
			else
				return e.setValue(value);
		} while (e != null);
		SimpleEntry<K, V> t = new SimpleEntry<>(key, value, parent);
		if (cmp < 0)
			parent.left = t;
		else
			parent.right = t;
		rebalanceAfterInsertion(t);
		size++;
		return null;
	}
	
	private void rebalanceAfterInsertion(SimpleEntry<K, V> e) {
		e.color = RED;
		if (e.parent == null) {
			e.color = BLACK;
			return;
		}
		if (color(e.parent) == RED) {
			SimpleEntry<K, V> dad = e.parent;
			SimpleEntry<K, V> uncle = null;
			SimpleEntry<K, V> grandpa = dad.parent;
			if (grandpa.left() == dad)
				uncle = grandpa.right();
			else
				uncle = grandpa.left();
			if (uncle == null || color(uncle) == BLACK) {
				if (grandpa.left() == dad) {
					if (dad.right() == e) {
						leftTurn(dad);
						e = dad;
						dad = e.parent;
						grandpa = dad.parent;
					}
					dad.color = BLACK;
					grandpa.color = RED;
					rightTurn(grandpa);
				} else if (grandpa.right() == dad) {
					if (dad.left() == e) {
						rightTurn(dad);
						e = dad;
						dad = e.parent;
						grandpa = dad.parent;
					}
					dad.color = BLACK;
					grandpa.color = RED;
					leftTurn(grandpa);
				}
			} else if (color(uncle) == RED) {
				dad.color = BLACK;
				uncle.color = BLACK;
				grandpa.color = RED;
				rebalanceAfterInsertion(grandpa);
			}
		}
	}

	public V remove(K key) {
		SimpleEntry<K, V> e = getEntry(key);
		if (e == null)
			return null;
		V oldValue = e.value;
		deleteEntry(e);
		return oldValue;
	}

	public void clear() {
		size = 0;
		root = null;
	}

	public boolean isEmpty() {
		return size == 0;
	}

	private void deleteEntry(SimpleEntry<K, V> e) {
		if (e.left != null && e.right != null) {
			SimpleEntry<K, V> succ = getSucc(e);
			e.key = succ.key;
			e.value = succ.value;
			e = succ;
		}
		SimpleEntry<K, V> child = null;
		if (e.left != null)
			child = e.left;
		else
			child = e.right;
		if (child != null) {
			child.parent = e.parent;
			if (e.parent == null)
				root = child;
			else if (e == e.parent.left)
				e.parent.left = child;
			else
				e.parent.right = child;
			e.left = e.right = e.parent = null;
			if (e.color == BLACK)
				rebalanceAfterDeletion(child);
		} else if (e.parent == null) {
			root = null;
		} else {
			if (e.color == BLACK)
				rebalanceAfterDeletion(e);
			if (e.parent != null) {
				if (e == e.parent.left)
					e.parent.left = null;
				else if (e == e.parent.right)
					e.parent.right = null;
				e.parent = null;
			}
		}
		size--;
	}

	private void rebalanceAfterDeletion(SimpleEntry<K, V> e) {
		while (e != root && color(e) == BLACK) {
			SimpleEntry<K, V> dad = e.parent;
			if (e == dad.left()) {
				SimpleEntry<K, V> child = dad.right();
				if (color(child) == RED) {
					child.color = BLACK;
					dad.color = RED;
					leftTurn(dad);
					child = dad.right();
				}
				if (color(child.left()) == BLACK && color(child.right()) == BLACK) {
					child.color = RED;
					e = dad;
					dad = parent(e);
				} else {
					if (color(child.right()) == BLACK) {
						child.left().color = BLACK;
						child.color = RED;
						rightTurn(child);
						child = dad.right();
					}
					child.color = color(dad);
					dad.color = BLACK;
					child.right().color = BLACK;
					leftTurn(dad);
					e = root;
				}
			} else {
				SimpleEntry<K, V> child = dad.left();
				if (color(child) == RED) {
					child.color = BLACK;
					dad.color = RED;
					rightTurn(dad);
					child = dad.left();
				}
				if (color(child.right()) == BLACK && color(child.left()) == BLACK) {
					child.color = RED;
					e = dad;
					dad = parent(e);
				} else {
					if (color(child.left()) == BLACK) {
						child.right().color = BLACK;
						child.color = RED;
						leftTurn(child);
						child = dad.left();
					}
					child.color = color(dad);
					dad.color = BLACK;
					child.left().color = BLACK;
					rightTurn(dad);
					e = root;
				}
			}
		}
		e.color = BLACK;
	}

	private SimpleEntry<K, V> getSucc(SimpleEntry<K, V> e) {
		if (e == null)
			return null;
		else if (e.right != null) {
			SimpleEntry<K, V> succ = e.right;
			while (succ.left != null)
				succ = succ.left;
			return succ;
		} else {
			SimpleEntry<K, V> succ = e.parent;
			SimpleEntry<K, V> child = e;
			while (succ != null && child == succ.right) {
				child = succ;
				succ = succ.parent;
			}
			return succ;
		}
	}

	public V get(K key) {
		SimpleEntry<K, V> e = getEntry(key);
		if(e != null)
			return e.value;
		else return null;
	}

	private SimpleEntry<K, V> getEntry(K key) {
		SimpleEntry<K, V> e = root;
		SimpleEntry<K, V> parent = null;
		int cmp = 0;
		do {
			parent = e;
			cmp = compare(key, e.key);			
			if (cmp < 0)
				e = e.left;
			else if (cmp > 0)
				e = e.right;
			else
				return e;
		} while (e != null);
		return null;
	}

	private void rightTurn(SimpleEntry<K, V> e) {
		if (e != null) {
			SimpleEntry<K, V> left = e.left;
			e.left = left.right;
			if (left.right != null)
				left.right.parent = e;
			left.parent = e.parent;
			if (e.parent == null)
				root = left;
			else if (e.parent.right == e)
				e.parent.right = left;
			else
				e.parent.left = left;
			left.right = e;
			e.parent = left;
		}
	}

	private void leftTurn(SimpleEntry<K, V> e) {
		if (e != null) {
			SimpleEntry<K, V> right = e.right;
			e.right = right.left;
			if (right.left != null)
				right.left.parent = e;
			right.parent = e.parent;
			if (e.parent == null)
				root = right;
			else if (e.parent.left == e)
				e.parent.left = right;
			else
				e.parent.right = right;
			right.left = e;
			e.parent = right;
		}
	}

	private SimpleEntry<K, V> getFirstEntry() {
		SimpleEntry<K, V> e = root;
		if (e != null)
			while (e.left != null)
				e = e.left;
		return e;
	}

	private boolean equalsOrNull(V v1, V v2) {
		if (v1 == null)
			return v2 == null;
		else
			return v1.equals(v2);
	}

	private int compare(K k1, K k2) {
		if (comparator == null) {
			return ((Comparable<K>) k1).compareTo((K) k2);
		} else {
			return comparator.compare(k1, k2);
		}
	}
	
	private boolean color(SimpleEntry<K, V> e) {
		if(e == null)
			return BLACK;
		else return e.color;
	}
	
	private SimpleEntry<K,V> parent(SimpleEntry<K,V> e) {
		if(e.parent != null)
			return e.parent;
		else return null;
	}

	private class EntryIterator<T> implements Iterator<MyMap.Entry<K, V>> {
		SimpleEntry<K, V> next;
		SimpleEntry<K, V> lastReturned;

		EntryIterator(SimpleEntry<K, V> first) {
			lastReturned = null;
			next = first;
		}

		public final boolean hasNext() {
			return next != null;
		}

		public MyMap.Entry<K, V> next() {
			SimpleEntry<K, V> e = next;
			if (e == null)
				throw new NoSuchElementException();
			next = getSucc(e);
			lastReturned = e;
			return e;
		}

		public void remove() {
			if (lastReturned == null)
				throw new IllegalStateException();
			if (lastReturned.left != null && lastReturned.right != null)
				next = lastReturned;
			deleteEntry(lastReturned);
			lastReturned = null;
		}
	}

	static class SimpleEntry<K, V> implements MyMap.Entry<K, V> {

		K key;
		V value;
		SimpleEntry<K, V> left = null;
		SimpleEntry<K, V> right = null;
		SimpleEntry<K, V> parent;
		boolean color = BLACK;

		SimpleEntry(K key, V value, SimpleEntry<K, V> parent) {
			this.key = key;
			this.value = value;
			this.parent = parent;
		}
		
		@Override
		public K getKey() {
			return key;
		}

		@Override
		public V getValue() {
			return value;
		}
		
		public SimpleEntry<K, V> left() {
			if(left == null)
				return null;
			else return left;
		}
		
		public SimpleEntry<K, V> right() {
			if(right == null)
				return null;
			else return right;
		}

		public V setValue(V value) {
			V old = this.value;
			this.value = value;
			return old;
		}

		public int hashCode() {
			return value.hashCode();
		}

		public boolean equals(Object obj) {
			return this.value.equals(obj);
		}

	}

}